import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import deque

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


# DRQN-GRU 网络
class DRQN_GRU(nn.Module):
    def __init__(self, width, height, action_size, sequence_length=10, gru_hidden_size=128):
        super(DRQN_GRU, self).__init__()
        self.sequence_length = sequence_length
        self.gru_hidden_size = gru_hidden_size
        self.action_size = action_size

        # 卷积层提取特征
        self.conv1 = nn.Conv2d(2, 16, kernel_size=3, stride=1, padding=1)  # 输入通道为2
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.flatten_size = 32 * width * height

        # GRU 层处理时间序列
        self.gru = nn.GRU(input_size=self.flatten_size, hidden_size=gru_hidden_size, batch_first=True)

        # 全连接层输出 Q 值
        self.fc1 = nn.Linear(gru_hidden_size, 128)
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(128, action_size)

    def forward(self, x, hidden=None):
        # x 的形状：[batch_size, sequence_length, channels, height, width]
        batch_size = x.size(0)

        # 处理卷积层
        conv_out = []
        for t in range(self.sequence_length):
            frame = x[:, t, :, :, :]
            frame = self.bn1(torch.relu(self.conv1(frame)))
            frame = self.bn2(torch.relu(self.conv2(frame)))
            frame = frame.view(batch_size, -1)  # 展平为 [batch_size, flatten_size]
            conv_out.append(frame)

        # 堆叠为序列输入 GRU
        conv_out = torch.stack(conv_out, dim=1)  # [batch_size, sequence_length, flatten_size]

        # GRU 前向传播
        gru_out, hidden = self.gru(conv_out, hidden)  # gru_out: [batch_size, sequence_length, gru_hidden_size]

        # 取最后一个时间步的输出
        gru_out = gru_out[:, -1, :]  # [batch_size, gru_hidden_size]

        # 全连接层
        x = torch.relu(self.fc1(gru_out))
        x = self.dropout(x)
        x = self.fc2(x)
        return x, hidden

    def init_hidden(self, batch_size):
        # 初始化 GRU 隐状态
        return torch.zeros(1, batch_size, self.gru_hidden_size).to(device)


# 序列经验回放缓冲区
class SequenceReplayBuffer:
    def __init__(self, capacity, sequence_length=10):
        self.buffer = deque(maxlen=capacity)
        self.sequence_length = sequence_length

    def push(self, sequence):
        # sequence 是一个包含 (state, action, reward, next_state, done) 的列表
        self.buffer.append(sequence)

    def sample(self, batch_size):
        # 随机采样 batch_size 个序列
        sequences = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = [], [], [], [], []
        for seq in sequences:
            state_seq, action_seq, reward_seq, next_state_seq, done_seq = zip(*seq)
            states.append(state_seq)
            actions.append(action_seq)
            rewards.append(reward_seq)
            next_states.append(next_state_seq)
            dones.append(done_seq)
        return (np.array(states), np.array(actions), np.array(rewards),
                np.array(next_states), np.array(dones))

    def __len__(self):
        return len(self.buffer)

    def clear(self):
        self.buffer.clear()


# DRQN-GRU 代理
class DQNAgent:
    def __init__(self, shape, action_size, sequence_length=40, gamma=0.99, epsilon=1,
                 epsilon_min=0.05, epsilon_decay=0.995, learning_rate=0.001,
                 batch_size=32, memory_size=100000):
        self.width = shape[1]
        self.height = shape[0]
        self.action_size = action_size
        self.sequence_length = sequence_length
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_min = epsilon_min
        self.epsilon_decay = epsilon_decay
        self.batch_size = batch_size
        self.memory = SequenceReplayBuffer(memory_size, sequence_length)

        # 初始化网络
        self.q_network = DRQN_GRU(shape[1], shape[0], action_size, sequence_length).to(device)
        self.target_network = DRQN_GRU(shape[1], shape[0], action_size, sequence_length).to(device)
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=learning_rate)

        # 用于测试时的隐状态
        self.hidden = None

    def reset_epsilon(self):
        self.epsilon = 1.0

    def reset_hidden(self):
        self.hidden = self.q_network.init_hidden(1)

    def is_over(self, action, env, q_net_args):
        if not env.take_actions(list(env.elements.keys())[0], int(action), sim_=True):
            q_net_args[:, action] = -10000
            return self.is_over(q_net_args.argmax().item(), env, q_net_args)
        return action

    def greedy_policy(self, bot, env):
        step_choice = [[0, -1], [0, 1], [-1, 0], [1, 0], [1, -1], [-1, -1], [1, 1], [-1, 1]]
        score = []
        quality_scores = np.array(env.Gaussian_mass_map)
        quality_scores_map = quality_scores / (np.max(quality_scores) + 1e-8)
        max_concentration_pos = np.unravel_index(np.argmax(quality_scores_map), quality_scores_map.shape)[::-1]
        for choice in step_choice:
            next_pos = np.array(bot.current_grid_pos) + choice
            if env.element_collection[bot.name].move(choice[0] * env.box_min_step[0], choice[1] * env.box_min_step[1],
                                                     sim_=True):
                distance = np.linalg.norm(np.array(max_concentration_pos) - next_pos).astype(
                    np.float16) + np.random.normal(0, 12)
                score.append(distance)
            else:
                score.append(float('inf'))  # 撞墙动作给无穷大距离
        # print(score, max_concentration_pos)
        # print(np.argmin(score), "选择动作")
        return np.argmin(score)

    def take_action(self, state_sequence, env, element_name):
        # state_sequence: [sequence_length, channels, height, width]
        state_tensor = torch.FloatTensor(state_sequence).unsqueeze(0).to(device)
        with torch.no_grad():
            if self.hidden is None:
                self.hidden = self.q_network.init_hidden(1)
            q_values, self.hidden = self.q_network(state_tensor, self.hidden)
            if random.random() < self.epsilon:
                # 随贪婪下降，随机过程的准确性会慢慢升高
                action = self.greedy_policy(env.element_collection[element_name], env)
            else:
                action = q_values.argmax().item()
        return action

    def store_experience(self, sequence):
        self.memory.push(sequence)

    def train(self):
        print(len(self.memory))
        if len(self.memory) < self.batch_size:
            return

        states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)
        states = torch.FloatTensor(states).to(device)  # [batch_size, seq_len, channels, h, w]
        actions = torch.LongTensor(actions).to(device)  # [batch_size, seq_len]
        rewards = torch.FloatTensor(rewards).to(device)
        next_states = torch.FloatTensor(next_states).to(device)
        dones = torch.FloatTensor(dones).to(device)

        # 初始化隐状态
        hidden = self.q_network.init_hidden(self.batch_size)
        target_hidden = self.target_network.init_hidden(self.batch_size)

        # 计算 Q 值
        q_values, _ = self.q_network(states, hidden)
        q_values = q_values.gather(1, actions[:, -1].unsqueeze(1)).squeeze(1)  # 取最后一步的动作

        # 计算目标 Q 值
        with torch.no_grad():
            next_q_values, _ = self.target_network(next_states, target_hidden)
            next_actions = next_q_values.argmax(1)
            target_q_values = rewards[:, -1] + (1 - dones[:, -1]) * self.gamma * \
                              next_q_values.gather(1, next_actions.unsqueeze(1)).squeeze(1)

        # 计算损失
        loss = nn.MSELoss()(q_values, target_q_values)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

    def update_target_network(self):
        self.target_network.load_state_dict(self.q_network.state_dict())

    def model_save(self):
        try:
            torch.save(self.q_network.state_dict(), 'drqn_eval.pkl')
            torch.save(self.target_network.state_dict(), 'drqn_target.pkl')
            print("successfully save")
        except Exception:
            print('save fail')

    def best_model_save(self):
        try:
            torch.save(self.q_network.state_dict(), 'best_drqn_eval.pkl')
            torch.save(self.target_network.state_dict(), 'best_drqn_target.pkl')
            print("successfully save")
        except Exception:
            print('save fail')

    def model_read(self):
        try:
            self.q_network.load_state_dict(torch.load('drqn_eval.pkl'))
            self.target_network.load_state_dict(torch.load('drqn_target.pkl'))
            print("successfully read")
        except Exception:
            print('read fail')
